#include "cfg.h"
#include <cassert>
using namespace cfg;
using namespace std;

CFG CFGGenerator::Generate(ir::Function& func)
{
    cfg.blocks.clear();
    block_map.clear();
    for (ir::Block& ir_block : func.blocks)
    {
        Block block;
        block.label = ir_block.label;
        block.instrs = ir_block.instrs;
        cfg.blocks.push_back(move(block));
    }
    cfg.entry_block = &cfg.blocks.front();
    for (Block& block : cfg.blocks) 
    { block_map[block.label] = &block; }
    for (Block& block : cfg.blocks)
    {
        assert(!block.instrs.empty()); 
        ir::Instruction& end_instr = block.instrs.back();
        switch (end_instr.op)
        {
            case ir::OP_BR: HandleBrBlock(block); break;
            case ir::OP_JMP: HandleJmpBlock(block); break;
            case ir::OP_RET: break;
            default: assert(false); break;
        }
    }
    return move(cfg);
}

void CFGGenerator::HandleBrBlock(Block& block)
{
    ir::Instruction& br_instr = block.instrs.back();
    assert(br_instr.srcs.size() == 3);
    assert(holds_alternative<string>(br_instr.srcs[1]));
    assert(holds_alternative<string>(br_instr.srcs[2]));
    string true_label = get<string>(br_instr.srcs[1]);
    string false_label = get<string>(br_instr.srcs[2]);
    assert(block_map.count(true_label));
    assert(block_map.count(false_label));
    AddEdge(&block, block_map[true_label]);
    AddEdge(&block, block_map[false_label]);
}

void CFGGenerator::HandleJmpBlock(Block& block)
{
    ir::Instruction& jmp_instr = block.instrs.back();
    assert(jmp_instr.srcs.size() == 1);
    assert(holds_alternative<string>(jmp_instr.srcs[0]));
    string label = get<string>(jmp_instr.srcs[0]);
    assert(block_map.count(label));
    AddEdge(&block, block_map[label]);
}

void CFGGenerator::AddEdge(Block* from, Block* to)
{
    from->succs.push_back(to);
    to->preds.push_back(from);
}